"""
Author: Morphlng
Date: 2024-03-15 11:19:41
LastEditTime: 2025-02-22 22:50:50
LastEditors: Morphlng
Description: Configuration for creating Carla adapter.
FilePath: /DrivingGym/src/driving_gym/simulation/adapter/carla/carla_config.py
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Callable, Dict, Optional, Tuple, ClassVar

import carla

import driving_gym.simulation.adapter.carla.utils as carla_util
from driving_gym.misc.util import NotProvided
from driving_gym.simulation.common.geom import Location, Rotation, Transform


@dataclass
class CarlaConfig:
    type: str = "carla"

    # Connection
    start_server: bool = False
    host: str = "localhost"
    port: Optional[int] = 2000
    timeout: float = 20.0
    offscreen: bool = False

    # Synchronization
    sync_mode: bool = True
    fixed_delta_seconds: Optional[float] = None
    command_in_batch: bool = False
    reload_epochs: int = -1

    # Traffic Manager
    tm_port: Optional[int] = None
    leading_distance: float = 5.0
    hybrid_physics_mode: bool = False
    hybrid_physics_radius: float = 50.0
    random_seed: Optional[int] = None

    # Spectator
    spectator_transform: Optional[Transform] = None
    spectator_follow_actor: Optional[str] = None

    def connection(
        self,
        host: str = "localhost",
        port: int = 2000,
        timeout: float = 20.0,
        start_server: bool = False,
        offscreen: bool = False,
    ):
        self.timeout = timeout
        self.offscreen = offscreen

        if start_server:
            self.start_server = True
            self.host = "localhost"
            self.port = None
        else:
            self.host = host
            self.port = port

        return self

    def spectator(self, transform=NotProvided, follow_actor: str = NotProvided):
        if transform is not NotProvided:
            self.spectator_transform = transform
        if follow_actor is not NotProvided:
            self.spectator_follow_actor = follow_actor
        return self

    def synchronization(
        self,
        sync_mode: bool = True,
        fixed_delta_seconds: float = 0.05,
        command_in_batch: bool = False,
        reload_epochs: int = -1,
    ):
        self.sync_mode = sync_mode
        self.fixed_delta_seconds = fixed_delta_seconds
        self.command_in_batch = command_in_batch
        self.reload_epochs = reload_epochs
        return self

    def traffic_manager(
        self,
        traffic_manager_port: int = None,
        leading_distance: float = 5.0,
        hybrid_physics_mode: bool = False,
        hybrid_physics_radius: float = 50.0,
        random_seed: int = None,
    ):
        self.tm_port = traffic_manager_port
        self.leading_distance = leading_distance
        self.hybrid_physics_mode = hybrid_physics_mode
        self.hybrid_physics_radius = hybrid_physics_radius
        self.random_seed = random_seed
        return self

    def to_dict(self):
        return asdict(self)


@dataclass
class CarlaActorConfig:
    actor_type: str = "vehicle"
    actor_blueprint: str = "vehicle.*"
    actor_spawn_point: Optional[Transform] = None
    actor_rolename: str = "npc"
    actor_attributes: Dict = field(default_factory=dict)
    autopilot: bool = False
    project_to_road: bool = True
    sensors: Dict = field(default_factory=dict)

    def blueprint(
        self,
        blueprint: str = "vehicle.*",
        type: str = "vehicle",
        rolename: str = "npc",
        attributes: dict = None,
        autopilot: bool = False,
    ):
        self.actor_type = type
        self.actor_blueprint = blueprint
        self.actor_rolename = rolename
        self.actor_attributes = attributes or {}
        self.autopilot = autopilot
        return self

    def spawning(self, spawn_point, project_to_road: bool = True):
        self.actor_spawn_point = spawn_point
        self.project_to_road = project_to_road
        return self

    def add_sensor(self, name: str, sensor_config: "dict | CarlaSensorConfig"):
        if not isinstance(sensor_config, dict):
            sensor_config = sensor_config.to_dict()

        self.sensors[name] = sensor_config
        return self

    @staticmethod
    def default_blueprint(actor_type: str):
        if "vehicle" in actor_type:
            bp = "vehicle.*"
        elif "pedestrian" in actor_type or "walker" in actor_type:
            bp = "walker.*"
        else:
            raise ValueError(f"You must provide a blueprint for: {actor_type}")
        return bp

    @staticmethod
    def parse_location(point, map=None):
        if point is None:
            return None
        elif map is not None:
            point = carla_util.parse_location(point, map)

        if isinstance(point, (Transform, carla.Transform)):
            return point
        elif isinstance(point, dict):
            return Transform(
                Location(
                    point.get("x", 0.0),
                    point.get("y", 0.0),
                    point.get("z", 0.0),
                ),
                Rotation(
                    point.get("pitch", 0.0),
                    point.get("yaw", 0.0),
                    point.get("roll", 0.0),
                ),
            )
        elif isinstance(point, (tuple, list)):
            return Transform(
                Location(*point[:3]),
                Rotation(*point[3:]),
            )
        elif isinstance(point, Location):
            return Transform(point, Rotation(0, 0, 0))
        else:
            raise ValueError(f"Unsupported spawn_point: {point}")

    def to_dict(self):
        d = asdict(self)
        for key in list(d.keys()):
            if key.startswith("actor_"):
                d[key[6:]] = d.pop(key)
        return d


@dataclass
class CarlaSensorConfig:
    _camera_transforms: ClassVar[list] = [
        {
            "x": 1.8,
            "y": 0.0,
            "z": 1.7,
            "pitch": 0.0,
            "yaw": 0.0,
            "roll": 0.0,
        },  # front camera for normal car
        {
            "x": 3.6,
            "y": 0.0,
            "z": 2.5,
            "pitch": 0.0,
            "yaw": 0.0,
            "roll": 0.0,
        },  # front camera for firetruck
        {
            "x": -5.5,
            "y": 0.0,
            "z": 2.8,
            "pitch": -15.0,
            "yaw": 0.0,
            "roll": 0.0,
        },  # back camera
        {
            "x": 1.8,
            "y": -0.5,
            "z": 1.7,
            "pitch": 0.0,
            "yaw": -90.0,
            "roll": 0.0,
        },  # left camera
        {
            "x": 1.8,
            "y": 0.5,
            "z": 1.7,
            "pitch": 0.0,
            "yaw": 90.0,
            "roll": 0.0,
        },  # right camera
    ]

    sensor_type: str = "rgb"
    sensor_blueprint: Optional[str] = None
    sensor_parent: Optional[str] = None
    sensor_pos_index: int = 0
    sensor_spawn_point: Optional[Transform] = None
    sensor_attachment_type: str = "rigid"
    sensor_attributes: Dict = field(default_factory=dict)

    relative_pos: bool = False
    relative_x: Tuple[float, float] = (0, 0)
    relative_y: Tuple[float, float] = (0, 0)
    relative_z: Tuple[float, float] = (0, 0)
    relative_pitch: float = 0
    relative_yaw: float = 0
    relative_roll: float = 0

    sensor_callback: Optional[Callable] = None

    def blueprint(
        self, type: str = "rgb", blueprint: str = None, attributes: dict = None
    ):
        self.sensor_type = type
        self.sensor_blueprint = blueprint
        self.sensor_attributes = attributes or {}
        return self

    def spawning(
        self,
        parent=None,
        spawn_point=None,
        attachment_type="rigid",
        pos_index=0,
    ):
        self.sensor_pos_index = pos_index
        self.sensor_spawn_point = CarlaActorConfig.parse_location(spawn_point)
        self.sensor_parent = parent
        self.sensor_attachment_type = attachment_type
        return self

    def position(
        self,
        relative_pos: bool = True,
        relative_x: tuple[float, float] = (0, 0),
        relative_y: tuple[float, float] = (0, 0),
        relative_z: tuple[float, float] = (0, 0),
        relative_pitch: float = 0,
        relative_yaw: float = 0,
        relative_roll: float = 0,
    ):
        """Position relative to the parent actor.

        Args:
            relative_pos (bool): Enable relative position. The following parameters will be used if True.
            relative_x (tuple): (coefficient, offset). Defaults to (1, 0).
            relative_y (tuple): (coefficient, offset). Defaults to (1, 0).
            relative_z (tuple): (coefficient, offset). Defaults to (1, 0).
            relative_pitch (int): Defaults to 0.
            relative_yaw (int): Defaults to 0.
            relative_roll (int): Defaults to 0.

        Note:
            - Take relative_x as an example, this will calculate as: `relative_x[0] * (bbox.extent.x + relative_x[1])`
            - Rotation will be set directly, e.g. `pitch = relative_pitch`
        """
        self.relative_pos = relative_pos
        self.relative_x = relative_x
        self.relative_y = relative_y
        self.relative_z = relative_z
        self.relative_pitch = relative_pitch
        self.relative_yaw = relative_yaw
        self.relative_roll = relative_roll
        return self

    def callback(self, func: Callable = None):
        """Register custom callback function for the sensor.

        Args:
            func (Callable, optional): The callback function. if None, will use the default callback.
        """
        self.sensor_callback = func
        return self

    @staticmethod
    def default_blueprint(sensor_type: str):
        sensor_type = sensor_type.lower()

        if "rgb" in sensor_type:
            bp = "sensor.camera.rgb"
        elif (
            "semantic" in sensor_type
            or "segment" in sensor_type
            or sensor_type == "semseg"
        ):
            bp = "sensor.camera.semantic_segmentation"
        elif "depth" in sensor_type:
            bp = "sensor.camera.depth"
        elif "lidar" in sensor_type:
            bp = "sensor.lidar.ray_cast"
        elif "radar" in sensor_type:
            bp = "sensor.other.radar"
        elif "imu" in sensor_type:
            bp = "sensor.other.imu"
        elif "gnss" in sensor_type:
            bp = "sensor.other.gnss"
        elif "colli" in sensor_type:
            bp = "sensor.other.collision"
        elif "lane" in sensor_type:
            bp = "sensor.other.lane_invasion"
        elif "obstacle" in sensor_type:
            bp = "sensor.other.obstacle"
        else:
            raise ValueError(f"You must provide a blueprint for: {sensor_type}")
        return bp

    def to_dict(self):
        d = asdict(self)
        for key in list(d.keys()):
            if key.startswith("sensor_"):
                d[key[7:]] = d.pop(key)
        return d
